package hooks

import (
	"MoSkeleton/framework/mo_redis"
	"MoSkeleton/framework/mo_redis/consts"
	"context"
	"github.com/redis/go-redis/v9"
	"github.com/samber/lo"
	"github.com/spf13/cast"
	"net"
	"strings"
)

// https://blog.csdn.net/qq_39272466/article/details/140825453
// 允许适配前缀的命令
var commandsWithPrefix = []string{
	"GET", "SET", "EXISTS", "DEL", "TYPE", "SETEX",
	"GETSET", "STRLEN", "DECR", "DECRBY", "APPEND",
	"GETRANGE", "SETRANGE", "GETBIT", "SETBIT",
	"RPUSH", "LPOP", "RPOP", "LLEN", "LRANGE",
	"SADD", "SREM", "SISMEMBER", "SMEMBERS", "SCARD",
	"HSET", "HMSET", "HGET", "HGETALL",
	"ZADD", "ZRANGE", "ZRANGEBYSCORE", "ZREVRANGEBYSCORE", "ZREM",
	"INCR", "INCRBY", "INCRBYFLOAT",
	"WATCH", "MULTI", "EXEC", "EXPIRE",
	//LIST
	"LINSERT", "LINDEX", "LREM", "LPUSH", "RPUSHX", "LSET", "LTRIM",
}

type KeyPrefix struct{}

func ShouldSkipPrefix(ctx context.Context) bool {
	value := ctx.Value(consts.Ctx_Skip_Prefix_Key)
	skip, ok := value.(bool)
	return ok && skip
}

/*
当创建网络连接时调用的hook
*/
func (k KeyPrefix) DialHook(next redis.DialHook) redis.DialHook {
	return func(ctx context.Context, network, addr string) (net.Conn, error) {
		return next(ctx, network, addr)
	}
}

/*
执行命令时调用的hook
*/
func (k KeyPrefix) ProcessHook(next redis.ProcessHook) redis.ProcessHook {
	return func(ctx context.Context, cmd redis.Cmder) error {
		if !ShouldSkipPrefix(ctx) {
			k.addKeyPrefix(ctx, cmd)
		}
		next(ctx, cmd)

		return nil
	}
}

/*
执行管道命令时调用的hook
*/
func (k KeyPrefix) ProcessPipelineHook(next redis.ProcessPipelineHook) redis.ProcessPipelineHook {
	return func(ctx context.Context, cmds []redis.Cmder) error {
		if !ShouldSkipPrefix(ctx) {
			for _, cmd := range cmds {
				k.addKeyPrefix(ctx, cmd)
			}
		}
		return next(ctx, cmds)
	}
}

func (k KeyPrefix) addKeyPrefix(ctx context.Context, cmd redis.Cmder) {
	args := cmd.Args()
	if len(args) <= 1 {
		return
	}
	cmdName := strings.ToUpper(cmd.FullName())
	switch cmdName {
	case "MGET", "DEL":
		for i := 1; i < len(args); i++ {
			if !ShouldSkipPrefix(ctx) {
				args[i] = k.combineKeyPrefix(args[i])
			}
		}
	case "MSET", "MSETNX":
		for i := 1; i < len(args); i += 2 {
			if !ShouldSkipPrefix(ctx) {
				args[i] = k.combineKeyPrefix(args[i])
			}
		}
	case "SCAN":
		if len(args) > 2 {
			for i := 1; i < len(args); i += 2 {
				if args[i] == "match" && i+1 < len(args) {
					args[i+1] = k.combineKeyPrefix(args[i+1])
					break
				}
			}
		}
	default:
		if lo.IndexOf[string](commandsWithPrefix, cmdName) != -1 {
			args[1] = k.combineKeyPrefix(args[1])
		}
	}
}

func (k KeyPrefix) combineKeyPrefix(arg interface{}) string {
	return mo_redis.RedisProperty.KeyPrefix + ":" + cast.ToString(arg)
}
